package cn.genmer.test.security.machinelearning.deeplearning4j.mnist.V2.strategy.concrete;

import cn.genmer.test.security.machinelearning.deeplearning4j.mnist.V1.MnistTrain;
import cn.genmer.test.security.machinelearning.deeplearning4j.mnist.V2.strategy.MnistModelStrategy;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.util.ModelSerializer;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;

public class DaeStrategy implements MnistModelStrategy {
    private Logger logger = LoggerFactory.getLogger(DaeStrategy.class);
  private MultiLayerNetwork model;
  public DaeStrategy() {
    // 加载预训练好的LeNet模型
    File locationToSave = new File(MnistTrain.BASE_PATH +"/dae_mnist_model.zip");
    System.out.println("预训练模型的路径：" + locationToSave.getAbsolutePath());
    try {
      model = ModelSerializer.restoreMultiLayerNetwork(locationToSave);
    } catch (IOException e) {
      e.printStackTrace();
    }
  }

    @Override
    public int predict(INDArray feature) {
        System.out.println("DAE: Denoising Autoencoder （去噪自编码器）");

        // 输入特征，模型预测
        INDArray output = model.output(feature);

        // 返回预测结果
        return Nd4j.argMax(output, 1).getInt(0);
    }

    @Override
    public INDArray loadImgAndGetFeature(BufferedImage image) {
        // Preprocess the image
        return preprocessImage(image);
    }

    @Override
    public INDArray loadImgAndGetFeature(String modelPath) {
        try {
            BufferedImage img = ImageIO.read(new File(modelPath));
            return loadImgAndGetFeature(img);
        } catch (Exception e){
            logger.error("【图像加载异常】", e);
        }
        return null;
    }


    private static INDArray preprocessImage(BufferedImage image) {

        // Resize the image to 28x28
        BufferedImage resizedImage = new BufferedImage(28, 28, BufferedImage.TYPE_BYTE_GRAY);
        resizedImage.getGraphics().drawImage(image, 0, 0, 28, 28, null);

        // Convert the image to INDArray
        INDArray array = Nd4j.create(1, 784);
        int[] pixels = new int[28 * 28];
        resizedImage.getRaster().getPixels(0, 0, 28, 28, pixels);
        for (int i = 0; i < pixels.length; i++) {
            array.putScalar(i, pixels[i]);
        }

        // Normalize the image
        double mean = array.meanNumber().doubleValue();
        double std = array.stdNumber().doubleValue();
        array.subi(mean).divi(std);

        return array;
    }
}